package 二叉树相关问题.中度题;

import java.util.LinkedList;
import java.util.List;

//不会
public class 二叉树中所有距离为k的节点 {
    List<Integer> ans;
    TreeNode target;
    int K;
    public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
        ans = new LinkedList<>();
        this.target=target;
        this.K=K;
        dfs(root);
        return ans;
    }

    public int dfs(TreeNode node) {
            if (node == null)
                return -1;
            else if (node == target) {
                subtree_add(node, 0);
                return 1;
            } else {
                int L = dfs(node.left), R = dfs(node.right);
                if (L != -1) {
                    if (L == K) ans.add(node.val);
                    subtree_add(node.right, L + 1);
                    return L + 1;
                } else if (R != -1) {
                    if (R == K) ans.add(node.val);
                    subtree_add(node.left, R + 1);
                    return R + 1;
                } else {
                    return -1;
                }
            }
        }

    //将距离节点K-dist的节点加入list中
    private void subtree_add(TreeNode root, int dist){
        if (root==null)
            return;
        if (dist==K)
            ans.add(root.val);
        else {
            subtree_add(root.left, dist + 1);
            subtree_add(root.right,dist+1);
        }
    }

}
